#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
双阶段 YOLO 检测节点

功能说明：
1. 订阅 /yolo_detect 话题接收控制命令（yolostart, yolostop等）
2. 订阅 /camera/image_raw 话题接收相机图像
3. 使用双阶段YOLO管道进行检测：
   - 第一阶段：在裁剪图像上检测锚点(数字、锥桶、仪表盘)
   - 第二阶段：对数字ROI进行二次YOLO检测
   - HSV颜色投票：对锥桶ROI进行颜色识别
4. 发布最终结果到/yolo_result话题
5. 发布状态信息到/yolo_status话题

话题接口：
- 订阅：/yolo_detect (std_msgs/String) - 控制命令
- 订阅：/camera/image_raw (sensor_msgs/Image) - 相机图像  
- 发布：/yolo_result (std_msgs/String) - 检测结果
- 发布：/yolo_status (std_msgs/String) - 状态信息

检测结果格式：color-number-state (如: red-three-up)
"""

import os
import sys
import gc
import time
import threading
from typing import Optional, Dict, Any

import rospy
import ctypes
import cv2
import numpy as np
from std_msgs.msg import String
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from collections import deque, Counter

# 添加yolo-jetson路径到sys.path以导入管道
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, "../../../../"))
yolo_jetson_path = os.path.join(project_root, "yolo-jetson", "detect")
if os.path.exists(yolo_jetson_path):
    sys.path.insert(0, yolo_jetson_path)

# 预加载 libgomp 以避免 "cannot allocate memory in static TLS block" 错误
try:
    ctypes.CDLL("libgomp.so.1", mode=ctypes.RTLD_GLOBAL)
except Exception as e:
    try:
        ctypes.CDLL("/lib/aarch64-linux-gnu/libgomp.so.1", mode=ctypes.RTLD_GLOBAL)
    except Exception as e2:
        rospy.logwarn(f"预加载 libgomp 失败: {e}; 备用路径失败: {e2}")

try:
    from double_yolo_pipeline import DoubleYoloPipeline
    rospy.loginfo("成功导入 DoubleYoloPipeline")
except ImportError as e:
    rospy.logerr(f"无法导入 DoubleYoloPipeline: {e}")
    rospy.logerr(f"检查路径: {yolo_jetson_path}")
    sys.exit(1)

# 引入GStreamerWrapper用于取流
try:
    from GStreamerWrapper import GStreamerWrapper
    rospy.loginfo("成功导入 GStreamerWrapper")
except Exception as e:
    rospy.logerr(f"无法导入 GStreamerWrapper: {e}")
    sys.exit(1)


class YoloDoubleStageNode:
    """双阶段YOLO检测ROS节点"""
    
    def __init__(self):
        rospy.init_node('yolo_double_stage_node', anonymous=True)
        
        # 参数配置
        self.model_path = rospy.get_param('~model_path', 
                                         '/home/jetson/robot_motion/yolo-jetson/weights/all_yolo_all_epochs_0811/best.pt')
        self.device = rospy.get_param('~device', None)  # None表示自动选择
        self.debug = rospy.get_param('~debug', True)
        self.save_dir = rospy.get_param('~save_dir', '/tmp/yolo_double_stage_debug')
        # 取流与投票参数
        self.gs_width = rospy.get_param('~gs_width', 640)
        self.gs_height = rospy.get_param('~gs_height', 360)
        self.queue_length = rospy.get_param('~queue_length', 7)
        self.frame_poll_interval = rospy.get_param('~frame_poll_interval', 0.02)  # GS拉流间隔
        # 预热与资源复用
        self.enable_warmup = rospy.get_param('~enable_warmup', True)
        self.warmup_iters = rospy.get_param('~warmup_iters', 3)
        self.warmup_timeout = rospy.get_param('~warmup_timeout', 2.0)
        self.warmup_use_gs = rospy.get_param('~warmup_use_gs', True)
        self.keep_gs_alive = rospy.get_param('~keep_gs_alive', True)
        self.keep_model_loaded = rospy.get_param('~keep_model_loaded', True)
        
        # 检测参数
        self.top_cut_ratio = rospy.get_param('~top_cut_ratio', 0.25)
        self.width_center_ratio = rospy.get_param('~width_center_ratio', 0.5)
        self.first_stage_conf = rospy.get_param('~first_stage_conf', 0.6)
        self.second_stage_conf = rospy.get_param('~second_stage_conf', 0.5)
        self.final_conf = rospy.get_param('~final_conf', 0.6)
        self.process_interval = rospy.get_param('~process_interval', 0.1)  # 处理间隔
        
        # 状态管理
        self.is_running = False
        self.bridge = CvBridge()
        self.pipeline: Optional[DoubleYoloPipeline] = None
        self.latest_image: Optional[np.ndarray] = None
        self.frame_count = 0
        self.last_process_time = 0
        self.processing_lock = threading.Lock()
        self.detection_thread: Optional[threading.Thread] = None
        self.stop_event = threading.Event()
        self.gs: Optional[GStreamerWrapper] = None
        
        # 结果缓存队列
        self.color_queue: deque = deque(maxlen=self.queue_length)
        self.number_queue: deque = deque(maxlen=self.queue_length)
        self.dashboard_queue: deque = deque(maxlen=self.queue_length)
        
        # 统计信息
        self.total_frames_processed = 0
        self.last_result = "None-None-None"
        self.last_processing_time = 0.0
        
        # ROS话题
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.status_pub = rospy.Publisher('/yolo_status', String, queue_size=10)
        
        self.command_sub = rospy.Subscriber('/yolo_detect', String, self.command_callback, queue_size=1)
        # 不再订阅ROS图像，改为GS取流
        
        rospy.loginfo("=== 双阶段YOLO检测节点初始化完成 ===")
        rospy.loginfo(f"模型路径: {self.model_path}")
        rospy.loginfo(f"设备: {self.device if self.device else '自动选择'}")
        rospy.loginfo(f"调试模式: {self.debug}")
        rospy.loginfo(f"检测参数: first_conf={self.first_stage_conf}, second_conf={self.second_stage_conf}, final_conf={self.final_conf}")
        rospy.loginfo("等待控制命令...")
        
        # 发布初始状态
        self.publish_status("节点已初始化，等待启动命令")

        # 启动后即预加载模型并进行一次（或若干次）预热推理，然后等待 yolostart
        try:
            if self.enable_warmup:
                rospy.loginfo("启动后进行模型预热以加速首帧...")
                # 若尚未加载模型，则先加载，后续 start_detection 复用
                if self.pipeline is None:
                    if not os.path.exists(self.model_path):
                        rospy.logwarn(f"预热跳过：模型文件不存在: {self.model_path}")
                    else:
                        self.pipeline = DoubleYoloPipeline(
                            model_path=self.model_path,
                            device=self.device,
                            debug=self.debug,
                            save_dir=self.save_dir
                        )
                # 如需用GS帧预热且要求保持取流线程常驻，则初始化GS
                if self.warmup_use_gs and self.keep_gs_alive and self.gs is None:
                    try:
                        self.publish_status("初始化GStreamer用于预热...")
                        self.gs = GStreamerWrapper(width=int(self.gs_width), height=int(self.gs_height))
                    except Exception as e:
                        rospy.logwarn(f"初始化GStreamer用于预热失败，将使用零帧：{e}")

                # 执行预热（使用零帧或GS帧）
                if self.pipeline is not None:
                    # 至少进行一次
                    if int(self.warmup_iters) <= 0:
                        self.warmup_iters = 1
                    self._warmup_models()
                    # 若不要求常驻，则关闭GS
                    if self.gs is not None and (not self.keep_gs_alive):
                        try:
                            self.gs.StopThread()
                        except Exception:
                            pass
                        self.gs = None
                    self.publish_status("模型已预热，等待 yolostart")
        except Exception as e:
            rospy.logwarn(f"启动预热失败: {e}")

    def publish_status(self, message: str):
        """发布状态信息"""
        try:
            status_msg = String()
            status_msg.data = f"[DoubleStage] {message}"
            self.status_pub.publish(status_msg)
        except Exception as e:
            rospy.logwarn(f"发布状态失败: {e}")

    def command_callback(self, msg: String):
        """处理控制命令"""
        command = msg.data.strip().lower()
        rospy.loginfo(f"收到命令: {command}")
        
        if command == "yolostart":
            self.start_detection()
        elif command == "yolostop":
            self.stop_detection()
        elif command == "yolostatus":
            self.report_status()
        elif command == "yoloreset":
            self.reset_system()
        else:
            rospy.logwarn(f"未知命令: {command}")
            self.publish_status(f"未知命令: {command}")

    def start_detection(self):
        """启动检测"""
        if self.is_running:
            rospy.loginfo("检测已在运行中")
            self.publish_status("检测已在运行中")
            return
            
        try:
            rospy.loginfo("正在初始化双阶段YOLO管道...")
            self.publish_status("正在初始化YOLO管道...")
            
            if not os.path.exists(self.model_path):
                error_msg = f"模型文件不存在: {self.model_path}"
                rospy.logerr(error_msg)
                self.publish_status(f"错误: {error_msg}")
                return
            
            # 初始化或复用YOLO管道
            if self.pipeline is None:
                self.pipeline = DoubleYoloPipeline(
                    model_path=self.model_path,
                    device=self.device,
                    debug=self.debug,
                    save_dir=self.save_dir
                )
            else:
                rospy.loginfo("复用已加载的YOLO模型与GPU缓存")
                self.publish_status("复用已加载的YOLO模型与GPU缓存")

            # 初始化或复用GStreamer取流
            if self.gs is None:
                self.publish_status("正在初始化GStreamer取流...")
                self.gs = GStreamerWrapper(width=int(self.gs_width), height=int(self.gs_height))
            else:
                # 若GS已存在但分辨率不同，必要时重建
                try:
                    if (hasattr(self.gs, 'width') and hasattr(self.gs, 'height') and
                        (self.gs.width != int(self.gs_width) or self.gs.height != int(self.gs_height))):
                        rospy.logwarn(
                            f"GS已在运行，当前分辨率 {self.gs.width}x{self.gs.height} 与期望 {self.gs_width}x{self.gs_height} 不一致，重新初始化GS")
                        try:
                            self.gs.StopThread()
                        except Exception:
                            pass
                        self.gs = GStreamerWrapper(width=int(self.gs_width), height=int(self.gs_height))
                    else:
                        rospy.loginfo("复用已在运行的GStreamer取流线程")
                        self.publish_status("复用GStreamer取流")
                except Exception:
                    # 无法确认状态则重建
                    try:
                        self.gs.StopThread()
                    except Exception:
                        pass
                    self.gs = GStreamerWrapper(width=int(self.gs_width), height=int(self.gs_height))
            
            self.is_running = True
            self.stop_event.clear()
            self.frame_count = 0
            self.total_frames_processed = 0
            self.last_process_time = time.time()
            # 清空历史队列
            self.color_queue.clear()
            self.number_queue.clear()
            self.dashboard_queue.clear()
            
            # 预热（在线程启动前执行，避免并发干扰）
            if self.enable_warmup:
                try:
                    self._warmup_models()
                except Exception as e:
                    rospy.logwarn(f"预热过程出现问题: {e}")

            # 启动后台识别线程
            self.detection_thread = threading.Thread(target=self._detection_loop, name="double_yolo_detection")
            self.detection_thread.daemon = True
            self.detection_thread.start()
            
            rospy.loginfo("✅ 双阶段YOLO检测已启动")
            self.publish_status("检测已启动，GS取流中...")
            
        except Exception as e:
            error_msg = f"启动检测失败: {e}"
            rospy.logerr(error_msg)
            self.publish_status(f"启动失败: {e}")
            self.is_running = False
            self.pipeline = None
            self.gs = None

    def stop_detection(self):
        """停止检测"""
        self.is_running = False
        self.stop_event.set()
        # 等待识别线程退出
        if self.detection_thread and self.detection_thread.is_alive():
            try:
                self.detection_thread.join(timeout=2.0)
            except Exception:
                pass
        self.detection_thread = None
        
        # 停止GStreamer（可选保留）
        try:
            if self.gs is not None:
                if self.keep_gs_alive:
                    rospy.loginfo("按配置保留GStreamer线程以加速下次启动")
                else:
                    try:
                        self.gs.StopThread()
                    except Exception:
                        pass
                    self.gs = None
        except Exception:
            pass
        
        # 释放模型（可选保留）
        if self.pipeline is not None:
            if self.keep_model_loaded:
                rospy.loginfo("按配置保留YOLO模型与GPU缓存以加速下次启动")
            else:
                try:
                    import torch
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                except Exception:
                    pass
                self.pipeline = None
                gc.collect()
        
        rospy.loginfo("🛑 双阶段YOLO检测已停止")
        self.publish_status(f"检测已停止，共处理 {self.total_frames_processed} 帧")

    def _get_warmup_frame(self) -> np.ndarray:
        """获取一帧用于预热，优先来自GS，超时后返回全零帧"""
        desired_w = int(self.gs_width)
        desired_h = int(self.gs_height)
        deadline = time.time() + float(self.warmup_timeout)
        if self.warmup_use_gs and self.gs is not None:
            while time.time() < deadline:
                frame = self.gs.GetFrame()
                if frame is not None:
                    return frame
                time.sleep(0.01)
        return np.zeros((desired_h, desired_w, 3), dtype=np.uint8)

    def _warmup_models(self):
        """在启动检测线程之前做若干次推理预热，触发模型与CUDA的lazy初始化"""
        iters = max(int(self.warmup_iters), 0)
        if iters <= 0 or self.pipeline is None:
            return
        self.publish_status(f"开始预热（{iters} 次）...")
        for i in range(iters):
            frame = self._get_warmup_frame()
            try:
                _ = self.pipeline.process_image(
                    image=frame,
                    top_cut_ratio=self.top_cut_ratio,
                    width_center_ratio=self.width_center_ratio,
                    first_stage_conf=self.first_stage_conf,
                    second_stage_conf=self.second_stage_conf,
                    final_conf=self.final_conf,
                    frame_idx=-(i + 1)
                )
            except Exception as e:
                rospy.logwarn(f"第 {i+1}/{iters} 次预热失败: {e}")
        self.publish_status("预热完成")

    def _value_is_valid(self, value: Optional[str]) -> bool:
        if value is None:
            return False
        v = str(value).strip().lower()
        if v in ("", "none", "unknown", "nan"):
            return False
        return True

    def _compute_majority(self, items: deque) -> str:
        if not items:
            return "None"
        counter = Counter(items)
        # 选出现次数最多的，若并列，选择最近出现的
        max_count = max(counter.values())
        candidates = [k for k, c in counter.items() if c == max_count]
        if len(candidates) == 1:
            return candidates[0]
        # 按从队尾到队首找第一个候选
        for v in reversed(items):
            if v in candidates:
                return v
        return candidates[0]

    def _detection_loop(self):
        """从GStreamer持续取流并识别，直到三类结果队列满后做多数投票并发布最终结果"""
        self.publish_status("识别线程已启动")
        while self.is_running and not self.stop_event.is_set():
            try:
                if self.gs is None:
                    time.sleep(0.1)
                    continue
                frame = self.gs.GetFrame()
                if frame is None:
                    time.sleep(self.frame_poll_interval)
                    continue
                self.latest_image = frame
                self.frame_count += 1
                
                # 控制处理节奏
                current_time = time.time()
                if current_time - self.last_process_time < self.process_interval:
                    time.sleep(self.frame_poll_interval)
                    continue
                
                # 处理当前帧
                start_time = time.time()
                result: Dict[str, Any] = self.pipeline.process_image(
                    image=self.latest_image.copy(),
                    top_cut_ratio=self.top_cut_ratio,
                    width_center_ratio=self.width_center_ratio,
                    first_stage_conf=self.first_stage_conf,
                    second_stage_conf=self.second_stage_conf,
                    final_conf=self.final_conf,
                    frame_idx=self.total_frames_processed + 1
                )
                processing_time = time.time() - start_time
                self.last_processing_time = processing_time
                self.last_process_time = time.time()
                self.total_frames_processed += 1
                
                if 'error' in result:
                    rospy.logwarn(f"处理错误: {result['error']}")
                    self.publish_status(f"处理错误: {result['error']}")
                    time.sleep(self.frame_poll_interval)
                    continue
                
                cone_color = result.get('cone_color', 'None')
                number = result.get('number', 'None')
                dashboard_state = result.get('dashboard_state', 'None')
                
                # 每帧输出识别结果
                rospy.loginfo(
                    f"[Frame {self.total_frames_processed}] 检测结果: "
                    f"{cone_color}-{number}-{dashboard_state} (耗时: {processing_time:.3f}s)"
                )

                # 仅将有效结果存入各自队列
                if self._value_is_valid(cone_color):
                    self.color_queue.append(str(cone_color))
                if self._value_is_valid(number):
                    self.number_queue.append(str(number))
                if self._value_is_valid(dashboard_state):
                    self.dashboard_queue.append(str(dashboard_state))
                
                # 日志
                rospy.loginfo(
                    f"[Frame {self.total_frames_processed}] 颜色Q/数字Q/仪表Q: "
                    f"{len(self.color_queue)}/{len(self.number_queue)}/{len(self.dashboard_queue)} "
                    f"(耗时: {processing_time:.3f}s)"
                )
                
                # 当三者队列都满时结算
                if (len(self.color_queue) == self.queue_length and
                    len(self.number_queue) == self.queue_length and
                    len(self.dashboard_queue) == self.queue_length):
                    final_color = self._compute_majority(self.color_queue)
                    final_number = self._compute_majority(self.number_queue)
                    final_dashboard = self._compute_majority(self.dashboard_queue)
                    final_result = f"{final_color}-{final_number}-{final_dashboard}"
                    self.last_result = final_result
                    
                    result_msg = String()
                    result_msg.data = final_result
                    self.result_pub.publish(result_msg)
                    rospy.loginfo(f"✅ 最终结果: {final_result}")
                    self.publish_status(f"最终结果发布: {final_result}")
                    
                    # 完成本次识别流程并停止
                    self.stop_detection()
                    break
                
                time.sleep(self.frame_poll_interval)
            except Exception as e:
                rospy.logerr(f"识别线程异常: {e}")
                self.publish_status(f"识别线程异常: {e}")
                time.sleep(0.1)

    def reset_system(self):
        """重置系统"""
        self.stop_detection()
        time.sleep(1)
        self.frame_count = 0
        self.total_frames_processed = 0
        self.last_result = "None-None-None"
        rospy.loginfo("🔄 系统已重置")
        self.publish_status("系统已重置")

    def report_status(self):
        """报告状态"""
        status = "运行中" if self.is_running else "已停止"
        uptime = time.time() - rospy.Time.now().to_sec() if rospy.Time.now().to_sec() > 0 else 0
        
        report = (f"状态: {status} | "
                 f"已处理: {self.total_frames_processed} 帧 | "
                 f"最后结果: {self.last_result} | "
                 f"处理耗时: {self.last_processing_time:.3f}s")
        
        rospy.loginfo(report)
        self.publish_status(report)

    # 以下方法保留以兼容，但当前未使用ROS图像流
    def image_callback(self, msg: Image):
        return

    def process_current_image(self):
        return

    def run(self):
        """运行节点"""
        rospy.loginfo("双阶段YOLO检测节点开始运行...")
        self.publish_status("节点运行中，等待控制命令")
        
        try:
            rospy.spin()
        except KeyboardInterrupt:
            rospy.loginfo("收到停止信号")
        finally:
            self.stop_detection()
            rospy.loginfo("双阶段YOLO检测节点已关闭")


if __name__ == '__main__':
    try:
        node = YoloDoubleStageNode()
        node.run()
    except rospy.ROSInterruptException:
        pass
    except Exception as e:
        rospy.logerr(f"节点启动失败: {e}")
        sys.exit(1)
